package matchbox;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

import basics.FeaturesSet;
import basics.Vector;
import basics.VectorMatrix;
import basics.VectorMatrixUtils;

public class NaiveBayes extends Classifier {

	private HashMap<Vector, HashMap<Integer, Double>> _pz_items = new HashMap<Vector, HashMap<Integer, Double>>();
	private HashMap<Vector, HashMap<Vector, HashMap<Double, Integer>>> _ratings = new HashMap<Vector, HashMap<Vector, HashMap<Double, Integer>>>();
	private HashMap<Integer, Double> _pz = new HashMap<Integer, Double>();
	private HashMap<Vector, HashMap<Vector, HashMap<Integer, HashMap<Double, Double>>>> _pr_z = new HashMap<Vector, HashMap<Vector, HashMap<Integer, HashMap<Double, Double>>>>();
	private int _k;

	private Vector _prev_user;
	private ArrayList<Vector> _items = new ArrayList<Vector>();
	private ArrayList<Double> _item_ratings = new ArrayList<Double>();
	private double _min;
	private double _max;

	@Override
	protected void trainUserItem(VectorMatrix u, VectorMatrix it, double r) {
		if (r == 0)
			return;

		if (u != _prev_user || _prev_user == null) {
			if (_prev_user != null) {
				performEM(_prev_user, _items, _item_ratings);
			}
			_prev_user = (Vector) u;
			_ratings.put(_prev_user, new HashMap<Vector, HashMap<Double, Integer>>());
			_pr_z.put(_prev_user, new HashMap<Vector, HashMap<Integer, HashMap<Double, Double>>>());
			_items.clear();
			_item_ratings.clear();
		}
		if (u == null || it == null)
			return;

		if (!_ratings.get(_prev_user).containsKey(it)) {
			_ratings.get(_prev_user).put((Vector) it, new HashMap<Double, Integer>());
		}

		if (_ratings.get(_prev_user).get(it).containsKey(r)) {
			_ratings.get(_prev_user).get(it).put(r, _ratings.get(_prev_user).get(it).get(r) + 1);
		} else {
			_ratings.get(_prev_user).get(it).put(r, 1);
		}

		if (!_pr_z.get(_prev_user).containsKey(it)) {
			_pr_z.get(_prev_user).put((Vector) it, new HashMap<Integer, HashMap<Double, Double>>());

			for (int z = 0; z < _k; z++) {
				_pr_z.get(_prev_user).get(it).put(z, new HashMap<Double, Double>());

				for (double rr = _min; rr <= _max; rr++) {
					_pr_z.get(_prev_user).get(it).get(z).put(rr, .1);
				}
			}
		}

		_items.add((Vector) it);
		_item_ratings.add(r);
	}

	private void performEM(Vector u, ArrayList<Vector> it, ArrayList<Double> r) {
		HashMap<Vector, HashMap<Integer, Double>> prev;
		int n = 0;
		do {
			prev = (HashMap<Vector, HashMap<Integer, Double>>) _pz_items.clone();
			computeE(u, it, r);
			computeM(u, it, r);
			n++;
		} while (!isConverged(prev) && n < 20);
	}

	private boolean isConverged(HashMap<Vector, HashMap<Integer, Double>> prev) {
		if (prev == null || prev.get(_prev_user) == null)
			return false;

		for (int i : prev.get(_prev_user).keySet()) {
			if (Math.abs(_pz_items.get(_prev_user).get(i) - prev.get(_prev_user).get(i)) > .2) {
				return false;
			}
		}

		return true;
	}

	private double getCachedRatingCount(Vector u, Vector it, double r) {
		int c = 0;
		try {
			c = _ratings.get(u).get(it).get(r);
		} catch (Exception ex) {
		}

		return c;
	}

	private double pr_zi(Vector uu, Vector it, double r, int z) {
		double p = 0, sum = 0.;
		for (double rr = _min; rr <= _max; rr++) {
			for (Vector u : _pz_items.keySet()) {
				p += _pz_items.get(u).get(z) * getCachedRatingCount(u, it, r);
			}
			p = p >= 0.01 ? p : 0.01;
			// _pr_z.get(uu).get(it).put(z, p);
			_pr_z.get(uu).get(it).get(z).put(rr, p);
			sum += p;
			p = 0;
		}

		for (double rr = _min; rr <= _max; rr++) {
			_pr_z.get(uu).get(it).get(z).put(rr, _pr_z.get(uu).get(it).get(z).get(rr) / sum);
		}
		return _pr_z.get(uu).get(it).get(z).get(r);
	}

	private void computeM(Vector u, ArrayList<Vector> it, ArrayList<Double> r) {
		updatePz();

		for (int i = 0; i < it.size(); i++) {
			// for (Double rr : r) {
			for (int z = 0; z < _k; z++) {
				// _pr_z.get(u).get(i).get(z).put(rr, );
				pr_zi(u, it.get(i), r.get(i), z);
			}
			// }
		}
	}

	private void updatePz() {
		double sum = 0;
		for (int z = 0; z < _k; z++) {
			double p = 0;
			for (Vector u : _pz_items.keySet()) {
				p += _pz_items.get(u).get(z);
			}
			_pz.put(z, p);
			sum += p;
		}
		for (int z = 0; z < _k; z++) {
			_pz.put(z, _pz.get(z) / sum);
		}
	}

	private void computeE(Vector u, List<Vector> it, List<Double> r) {
		double p = 0., sum = 0;
		for (int z = 0; z < _k; z++) {
			p = _pz.get(z);
			for (int i = 0; i < it.size(); i++) {
				p *= _pr_z.get(u).get(it.get(i)).get(z).get(r.get(i));
			}
			p = p == 0 ? p = 1E-6 : p;
			// p = p < 0.01 ? .01 : p;
			if (!_pz_items.containsKey(u)) {
				_pz_items.put(u, new HashMap<Integer, Double>());
			}
			_pz_items.get(u).put(z, p);
			sum += p;
		}

		for (int z = 0; z < _k; z++) {
			_pz_items.get(u).put(z, _pz_items.get(u).get(z) / sum);
		}
	}

	// private double pr_z(Vector it, double r, int z) {
	// double p = 0, s = 0;
	// for (Vector u : _pz_items.keySet()) {
	// if (_ratings.get(u).get(it) == r) {
	// p += _pz_items.get(u).get(z);
	// }
	// s += _pz_items.get(u).get(z);
	// }
	//
	// return p / s;
	// }
	//
	// private double pr_z(List<Vector> it, List<Double> r, int z) {
	// double p = 1;
	// // for (Vector i : it) {
	// for (double rr : r) { // _pr_z.keySet()
	// // p *= pr_z(i, rr, z);
	// p *= _pr_z.get(rr).get(z);
	// }
	// // }
	// return p > .0001 ? p : 0.0001;
	// }

	@Override
	public double infer(VectorMatrix u, VectorMatrix it) {
		trainUserItem(null, null, -1);

		double max_r = 3, val = 0;
		for (double rr = _min; rr <= _max; rr++) {
			for (int z = 0; z < _k; z++) {
				if (_pr_z.get(u).get(it).get(z).get(rr) > val) {
					val = _pr_z.get(u).get(it).get(z).get(rr);
					max_r = rr;
				}
			}
		}

		return max_r;
	}

	@Override
	public void train(FeaturesSet users, FeaturesSet items, FeaturesSet ratings, int k) {
		_k = k;

		_min = ratings.min(VectorMatrixUtils.OperationOpt.NonZero);
		_max = ratings.max(VectorMatrixUtils.OperationOpt.NonZero);

		for (double i = _min; i <= _max; i++) {
			// _pr_z.put(i, new HashMap<Integer, Double>());

			for (int z = 0; z < _k; z++) {
				_pz.put(z, 1. / (double) _k);
				// _pr_z.get(i).put(z, .1);
			}
		}

		for (int i = 0; i < users.size(); i++) {
			for (int j = 0; j < items.size(); j++) {
				trainUserItem(users.get(i), items.get(j), ratings.get(i).get(j));
			}
		}
	}
}
